import os
from typing import Iterable

import numpy as np
from zkl_datasets import ChunkedBinaryDatasetWriter, DatasetWriter, MappedDatasetWriter, RandomDispatchingDirGroupedDatasetsWriter

from .vocab import SimpleVocabWithCount


class MultiSplitsTokenizedDatasetWriter(RandomDispatchingDirGroupedDatasetsWriter[np.ndarray, DatasetWriter]):
    def __init__(self, path: os.PathLike | str, ratios: dict[str, float] | Iterable[tuple[str, float]], vocab: SimpleVocabWithCount):
        def child_writer_factory(path: os.PathLike | str) -> DatasetWriter[np.ndarray]:
            child_writer = ChunkedBinaryDatasetWriter(path)
            child_writer = MappedDatasetWriter(child_writer, lambda sample: sample.tobytes())
            return child_writer

        super().__init__(path, ratios, child_writer_factory)
        vocab.save_to_file(os.path.join(self._path, "vocab.csv"))
